這張提到 JAX 的投影片來自 State of AI Report 2021 (https://www.stateof.ai/2021) 這份報告的第 74 頁，它的標題是：「Here comes a new framework challenger: JAX」，它將 JAX 定位為「AI 框架的新進挑戰者」。也就是說，除了 TensorFlow 和 Pytorch 之外, State of AI Report 認為 JAX 即將 (或者是已經 ?) 成為我們第三個 AI 框架的選項！（註：現行的 AI 框架當然不限於 TensorFlow 及 Pytorch，只是這幾年來，老頭的工作裏，都是使用 TensorFlow 或 Pytorch，已經很少用到其他的框架。）

在此之前，老頭雖曾耳聞 JAX ，但是我一直認為 JAX 只是一個 Python 函式庫，提供有效率的矩陣及微分運算，從來沒有把它跟 AI 框架聯想在一起！這份報告激起了老頭對 JAX 莫大的興趣，覺得我應該花點時間，把 JAX 好好的認識一番，進而判斷我們是不是真的可以選擇 JAX 做為框架，來開發未來的 AI 產品。

讀者可能想問：「老頭，那你有答案了嗎？」有的，但是這個答案不是三言兩語可以說清楚的。只能簡單的講，「適狀況而定！」老頭希望藉著這一系列的發文，讀者們能夠更深入的了解 JAX，進而有能力針對你自己的狀況，找出你自己的答案。如果你真的很急，想要馬上得到一些建議，那麼你可以參考這篇文章：

Why You Should (or Shouldn't) be Using Google's JAX in 2022
(https://www.assemblyai.com/blog/why-you-should-or-shouldnt-be-using-jax-in-2022/)

雖然老頭並不完全同意它的看法，但是它提供了許多中肯的建議，很值得參考。老頭會接下來的發文，在適當的時候，也會引用它裏面諸多精采的內容。

再回到這張投影片，我把它的重點摘要出來，讓大家對 JAX 有一個概略的印象：

• JAX 是由 Google 在 2019 年提出的。（註：其實在 2018 年，Google 就提出了論文 )
• JAX 是結合了之前的兩個技術而成的，它們是 Autograd 和 XLA。
• JAX 模仿 Numpy。
• JAX 的重要功能包括 easy vectorization (向量化), parallelization (平行處理化), just-in-time compilaton (即時編譯)。
• JAX 生態系正在迅速的成長，許多擴增的函式庫被發表，包括 Flax, Haiku, Optax, Rlax, FedJAX 等等。
• 預期 JAX 很快的就會由學界走到業界。

看不太懂沒有關係，接下來就讓老頭一步步地帶大家走進 JAX 的世界裏。